import argparse
import torch
import yaml
import os
from yolov5 import train as yolov5_train
from yolov5.models.yolo import Model
from utils.anchors import calculate_adaptive_anchors
from data.augmentation import MosaicAugmentation, CutMixAugmentation

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='data/config.yaml', help='dataset config file path')
    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
    parser.add_argument('--batch-size', type=int, default=16, help='batch size')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--adaptive-anchors', action='store_true', help='use adaptive anchor strategy')
    parser.add_argument('--mosaic', type=float, default=1.0, help='mosaic augmentation probability')
    parser.add_argument('--cutmix', type=float, default=1.0, help='cutmix augmentation probability')
    return parser.parse_args()

def main(opt):
    # Load dataset configuration
    with open(opt.data) as f:
        data_config = yaml.safe_load(f)

    # Calculate adaptive anchors if enabled
    if opt.adaptive_anchors:
        print("Calculating adaptive anchors...")
        anchors = calculate_adaptive_anchors(data_config['train'])
        # Update anchors in configuration
        with open(opt.data, 'r') as f:
            config_content = f.read()
        config_content = config_content.replace('anchors:', f'anchors: {anchors}')
        with open(opt.data, 'w') as f:
            f.write(config_content)

    # Load YOLOv5 model
    model = Model(cfg='models/yolov5s.yaml', ch=3, nc=data_config['nc'])
    if opt.weights.endswith('.pt'):
        model.load_state_dict(torch.load(opt.weights)['model'].state_dict())

    # Initialize custom data augmentations
    transforms = []
    if opt.mosaic > 0:
        transforms.append(MosaicAugmentation(probability=opt.mosaic))
    if opt.cutmix > 0:
        transforms.append(CutMixAugmentation(probability=opt.cutmix))

    # Start training using YOLOv5's training function with custom parameters
    opt.rect = False  # Disable rectangular training for custom augmentation compatibility
    yolov5_train.run(**vars(opt))

if __name__ == '__main__':
    opt = parse_opt()
    main(opt)